import datetime as dt
import pandas as pd
from datetime import datetime, timedelta


def initialize(context):
    """
    策略初始化函数，用于设置全局交易参数和每日操作计划。
    
    参数:
    - context: 策略运行上下文，包含全局变量和函数。
    
    此函数中设置了交易的各种参数，包括交易量比例、限制模式、佣金和滑点，以及每日买入和卖出的操作计划。
    """
    # 设置交易量比例为1，意味着全额交易
    set_volume_ratio(volume_ratio=1)
    # 设置限制模式为无限制，意味着在交易中没有位置或头寸的限制
    set_limit_mode(limit_mode='UNLIMITED')
    # 设置佣金比例为万分之一，最低佣金为5元，适用于股票交易
    set_commission(commission_ratio=1 / 10000, min_commission=5.0, type="STOCK")
    # 设置固定滑点为万分之三，模拟交易中的价格不利变动
    set_fixed_slippage(fixedslippage=3 / 10000)

    # 每日09:31运行买入操作，根据策略逻辑买入股票
    run_daily(context, buy, time='09:31')
    # 每日10:00、11:25和14:50运行卖出操作，根据策略逻辑卖出股票
    run_daily(context, sell, time='10:00')
    run_daily(context, sell, time='11:25')
    run_daily(context, sell, time='14:50')
    # 新增：初始化全局变量，用于存储当天买入的股票
    g.today_bought_stocks = []  # 当天买入的股票列表
    g.historical_bought_stocks = []  # 历史买入的股票列表（可选）


def before_trading_start(context, data):
    """
    在每个交易日开始前调用此函数，以准备当天所需的股票列表。

    :param context: 上下文对象，存储全局变量和状态。
    :param data: 数据对象，包含市场数据。
    """
    # 获取最近三个交易日的日期
    # 获取最近一个交易日
    date_1 = get_trading_day(-1)
    date_2 = get_trading_day(-2)
    date_3 = get_trading_day(-3)

    # 准备初始股票列表(所有A股股票 - 停牌 - 退市 - ST - 暂停上市 - 新股)
    initial_list = prepare_stock_list(date_1)

    # 获取在date_1达到历史高点的股票列表
    hl0_list = get_hl_stock(initial_list, date_1)
    # 获取在date_2和date_3达到历史高点的股票列表
    hl1_list = get_ever_hl_stock(initial_list, date_2)
    hl2_list = get_ever_hl_stock(initial_list, date_3)

    # 合并date_2和date_3的高点股票列表，用于后续的差集操作
    elements_to_remove = set(hl1_list + hl2_list)

    # 计算date_1的高点股票中，哪些在date_2和date_3没有达到高点，即为g.gap_up列表
    g.gap_up = [stock for stock in hl0_list if stock not in elements_to_remove]

    # 计算date_1的高点股票中，哪些在date_2没有达到高点，即为g.gap_down列表
    g.gap_down = [s for s in hl0_list if s not in hl1_list]

    # 获取在date_1达到历史高点的股票列表，参数True表示包含当日
    h1_list = get_ever_hl_stock(initial_list, date_1, True)

    # 获取在date_2达到历史高点的股票列表，用于后续的差集操作
    elements_to_remove = get_hl_stock(initial_list, date_2)

    # 计算在date_1达到历史高点但date_2没有达到高点的股票列表，即为g.reversal列表
    g.reversal = [stock for stock in h1_list if stock not in elements_to_remove]


def buy(context):
    qualified_stocks = []
    gk_stocks = []
    dk_stocks = []
    rzq_stocks = []

    # 用于记录未通过的股票及其原因
    disqualify_details = {
        'gap_up': {},
        'gap_down': {},
        'reversal': {}
    }

    # 获取所有股票名称（一次性获取提高效率）
    all_stocks = list(g.gap_up) + list(g.gap_down) + list(g.reversal)
    stock_names = get_stock_name(all_stocks) if all_stocks else {}

    # 高开股票筛选
    for s in g.gap_up:
        disqualify_reasons = []

        # 条件一：均价，金额，市值，换手率
        prev_day_data = get_history(count=1, frequency='1d', field=['close', 'volume', 'money'], security_list=s)
        avg_price_increase_value = prev_day_data['money'][0] / prev_day_data['volume'][0] / prev_day_data['close'][
            0] * 1.1 - 1
        if avg_price_increase_value < 0.07:
            disqualify_reasons.append("均价涨幅<7%")
        if prev_day_data['money'][0] < 5.5e8:
            disqualify_reasons.append("成交额<5.5亿")
        if prev_day_data['money'][0] > 20e8:
            disqualify_reasons.append("成交额>20亿")

        # 市值条件
        turnover_ratio_data = get_fundamentals(s, table='valuation', fields=['total_value', 'float_value'])
        if turnover_ratio_data.empty:
            disqualify_reasons.append("无市值数据")
        else:
            if turnover_ratio_data['total_value'][0] < 7e9:
                disqualify_reasons.append("总市值<70亿")
            if turnover_ratio_data['float_value'][0] > 5.2e10:
                disqualify_reasons.append("流通市值>520亿")

        # 左压条件
        if rise_low_volume(s, context):
            disqualify_reasons.append("左压不满足")

        # 高开条件
        auction_data = get_trend_data(stocks=s)
        if not auction_data or auction_data[s] is None:
            disqualify_reasons.append("无竞价数据")
        else:
            if auction_data[s]['business_amount'] / prev_day_data['volume'][-1] < 0.03:
                disqualify_reasons.append("竞价量<3%")

            current_data = get_history(count=1, frequency='1d', field=['high_limit'], security_list=s, include=True)
            current_ratio = auction_data[s]['hq_px'] / (current_data['high_limit'][0] / 1.1)
            if current_ratio <= 1:
                disqualify_reasons.append("非高开")
            if current_ratio >= 1.06:
                disqualify_reasons.append("高开>6%")

        if disqualify_reasons:
            for reason in disqualify_reasons:
                if reason not in disqualify_details['gap_up']:
                    disqualify_details['gap_up'][reason] = []
                disqualify_details['gap_up'][reason].append(f"{s}({stock_names.get(s, '未知')})")
        else:
            gk_stocks.append(s)
            qualified_stocks.append(s)

    # 低开股票筛选
    date = context.previous_date.strftime("%Y-%m-%d")
    if g.gap_down:
        stock_list = g.gap_down
        rpd = get_relative_position_df(stock_list, date, 60)
        rpd = rpd[rpd['rp'] <= 0.5]
        stock_list = list(rpd.index)

        if stock_list:
            df = get_history(count=1, frequency='1d', field=['preclose', 'open'], security_list=stock_list,
                             include=True)
            df['open_pct'] = df['open'] / df['preclose']
            df = df[(0.955 <= df['open_pct']) & (df['open_pct'] <= 0.97)]
            stock_list = df['code'].tolist()

        for s in stock_list:
            disqualify_reasons = []
            prev_day_data = get_history(count=1, frequency='1d', field=['close', 'volume', 'money'], security_list=s)
            if prev_day_data['money'][0] < 1e8:
                disqualify_reasons.append("成交额<1亿")

            if disqualify_reasons:
                for reason in disqualify_reasons:
                    if reason not in disqualify_details['gap_down']:
                        disqualify_details['gap_down'][reason] = []
                    disqualify_details['gap_down'][reason].append(f"{s}({stock_names.get(s, '未知')})")
            else:
                dk_stocks.append(s)
                qualified_stocks.append(s)

    # 弱转强股票筛选
    for s in g.reversal:
        disqualify_reasons = []

        # 三日涨幅
        prev_day_data = get_history(count=4, frequency='1d', field=['open', 'close', 'volume', 'money'],
                                    security_list=s)
        increase_ratio = (prev_day_data['close'][-1] - prev_day_data['close'][0]) / prev_day_data['close'][0]
        if increase_ratio > 0.28:
            disqualify_reasons.append("三日涨幅>28%")

        # 前一日收盘表现
        open_close_ratio = (prev_day_data['close'][-1] - prev_day_data['open'][-1]) / prev_day_data['open'][-1]
        if open_close_ratio < -0.05:
            disqualify_reasons.append("前日收盘<-5%")

        # 均价和金额
        avg_price_increase_value = prev_day_data['money'][-1] / prev_day_data['volume'][-1] / prev_day_data['close'][
            -1] - 1
        if avg_price_increase_value < -0.04:
            disqualify_reasons.append("均价涨幅<-4%")
        if prev_day_data['money'][-1] < 3e8:
            disqualify_reasons.append("成交额<3亿")
        if prev_day_data['money'][-1] > 19e8:
            disqualify_reasons.append("成交额>19亿")

        # 市值条件
        turnover_ratio_data = get_fundamentals(s, table='valuation', fields=['total_value', 'float_value'])
        if turnover_ratio_data.empty:
            disqualify_reasons.append("无市值数据")
        else:
            if turnover_ratio_data['total_value'][0] < 7e9:
                disqualify_reasons.append("总市值<70亿")
            if turnover_ratio_data['float_value'][0] > 5.2e10:
                disqualify_reasons.append("流通市值>520亿")

        # 左压条件
        if rise_low_volume(s, context):
            disqualify_reasons.append("左压不满足")

        # 高开条件
        auction_data = get_trend_data(stocks=s)
        if not auction_data or auction_data[s] is None:
            disqualify_reasons.append("无竞价数据")
        else:
            if auction_data[s]['business_amount'] / prev_day_data['volume'][-1] < 0.03:
                disqualify_reasons.append("竞价量<3%")

            current_data = get_history(count=1, frequency='1d', field=['high_limit'], security_list=s, include=True)
            current_ratio = auction_data[s]['hq_px'] / (current_data['high_limit'][0] / 1.1)
            if current_ratio <= 0.98:
                disqualify_reasons.append("开盘<-2%")
            if current_ratio >= 1.09:
                disqualify_reasons.append("开盘>9%")

        if disqualify_reasons:
            for reason in disqualify_reasons:
                if reason not in disqualify_details['reversal']:
                    disqualify_details['reversal'][reason] = []
                disqualify_details['reversal'][reason].append(f"{s}({stock_names.get(s, '未知')})")
        else:
            rzq_stocks.append(s)
            qualified_stocks.append(s)

    # 精简版汇总日志
    log.info('\n========== 最终选股结果 ==========')

    # 买入的股票
    if qualified_stocks:
        log.info('【买入股票】')
        for category, stocks in [('一进二(高开)', gk_stocks), ('首板低开', dk_stocks), ('弱转强', rzq_stocks)]:
            if stocks:
                log.info(f'{category}:')
                for s in stocks:
                    log.info(f'  {s}({stock_names.get(s, "未知")})')
    else:
        log.info('今日无符合买入条件的股票')

    # 未通过的股票按原因分类
    log.info('\n【未通过筛选的股票】')
    for category in ['gap_up', 'gap_down', 'reversal']:
        chinese_name = {
            'gap_up': '一进二(高开)',
            'gap_down': '首板低开',
            'reversal': '弱转强'
        }[category]

        if disqualify_details[category]:
            log.info(f'{chinese_name}类:')
            for reason, stocks in disqualify_details[category].items():
                log.info(f'  {reason}:')
                # 每行显示5只股票，避免单行过长
                for i in range(0, len(stocks), 5):
                    log.info('    ' + ', '.join(stocks[i:i + 5]))
        else:
            log.info(f'{chinese_name}类: 无未通过筛选的股票')

    # 执行买入
    if len(qualified_stocks) != 0 and context.portfolio.cash / context.portfolio.portfolio_value > 0.3:
        value = context.portfolio.cash / len(qualified_stocks)
        for s in qualified_stocks:
            posInfo = get_position(s)
            amount = value / posInfo.last_sale_price // 100 * 100
            if amount >= 100:
                if is_trade():
                    order_market(s, amount, 2)
                else:
                    order_target(s, amount)
            log.info(f'执行买入 {s}({stock_names.get(s, "未知")}), 数量: {amount}')

    # 执行买入前，清空当天买入列表（避免重复）
    g.today_bought_stocks = []

    # 执行买入并记录
    if len(qualified_stocks) != 0 and context.portfolio.cash / context.portfolio.portfolio_value > 0.3:
        value = context.portfolio.cash / len(qualified_stocks)
        for s in qualified_stocks:
            posInfo = get_position(s)
            amount = value / posInfo.last_sale_price // 100 * 100
            if amount >= 100:
                if is_trade():
                    order_market(s, amount, 2)
                else:
                    order_target(s, amount)
                # 记录实际买入的股票（确保成交后再记录）
                g.today_bought_stocks.append(s)
                log.info(f'买入 {s}({get_stock_name(s)[s]}), 数量: {amount}')

    # 可选：将当天买入股票加入历史记录
    g.historical_bought_stocks.extend(g.today_bought_stocks)


def sell(context):
    if str(context.blotter.current_dt)[-8:] == '10:00:00':
        for s in list(context.portfolio.positions):
            current_data = get_history(count=1, frequency='1d', field=['close', 'open', 'preclose'], security_list=s, include=True)

            close = current_data['close'][-1]
            open = current_data['open'][-1]
            preclose = current_data['preclose'][-1]

            sell_flag = close < open or close < preclose

            enable_amount = context.portfolio.positions[s].enable_amount
            last_sale_price = context.portfolio.positions[s].last_sale_price
            cost_basis = context.portfolio.positions[s].cost_basis

            if (sell_flag and enable_amount > 0):
                if is_trade():
                    log.info("当前运行场景：交易")
                    order_market(s, enable_amount * -1, 2)
                else:
                    log.info("当前运行场景：回测")
                    order_target(s, enable_amount * -1)
                log.info('半小时不见红卖出 {} : {}'.format(s, get_stock_name(s)[s]))
                print('———————————————————————————————————')

    # 根据时间执行不同的卖出策略
    if str(context.blotter.current_dt)[-8:] == '11:25:00':
        for s in list(context.portfolio.positions):
            current_data = get_history(count=1, frequency='1d', field=['close', 'high_limit'], security_list=s, include=True)

            high_limit = current_data['high_limit'][-1]
            enable_amount = context.portfolio.positions[s].enable_amount
            last_sale_price = context.portfolio.positions[s].last_sale_price
            cost_basis = context.portfolio.positions[s].cost_basis

            # 早盘没涨停的时候就卖出
            if ((enable_amount > 0) and (last_sale_price < high_limit) and (last_sale_price > cost_basis)):
                if is_trade():
                    log.info("当前运行场景：交易")
                    order_market(s, enable_amount * -1, 2)
                else:
                    log.info("当前运行场景：回测")
                    order_target(s, enable_amount * -1)
                log.info('止盈卖出 {} : {}'.format(s, get_stock_name(s)[s]))
                print('———————————————————————————————————')

    if str(context.current_dt)[-8:] == '14:50:00':
        for s in list(context.portfolio.positions):

            close_data2 = get_history(count=5, frequency='1d', field=['close', 'high_limit'], security_list=s, include=True)
            MA5 = close_data2['close'].mean()

            high_limit = close_data2['high_limit'][-1]
            enable_amount = context.portfolio.positions[s].enable_amount
            last_sale_price = context.portfolio.positions[s].last_sale_price
            cost_basis = context.portfolio.positions[s].cost_basis

            if ((enable_amount > 0) and (last_sale_price < high_limit) and (last_sale_price > cost_basis)):
                if is_trade():
                    log.info("当前运行场景：交易")
                    order_market(s, enable_amount * -1, 2)
                else:
                    log.info("当前运行场景：回测")
                    order_target(s, enable_amount * -1)
                log.info('止盈卖出 {} : {}'.format(s, get_stock_name(s)[s]))
                print('———————————————————————————————————')

            elif ((enable_amount > 0) and (last_sale_price < MA5)):
                # enable_amount可卖出的仓位
                if is_trade():
                    log.info("当前运行场景：交易")
                    order_market(s, enable_amount * -1, 2)
                else:
                    log.info("当前运行场景：回测 ")
                    order_target(s, enable_amount * -1)
                log.info('止损卖出 {} : {} 数量：{}'.format(s, get_stock_name(s)[s], enable_amount))
                print('———————————————————————————————————')


def get_hl_stock(stock_list, date):
    """
    获取指定日期的涨停股票列表。

    该函数首先将给定日期转换为特定格式的字符串，然后调用get_price函数获取指定股票列表的历史价格数据。
    通过筛选条件，找出当天涨停且开盘的股票，并返回其代码列表。

    参数:
    stock_list: 股票代码列表，用于指定需要查询的股票。
    date: datetime对象，表示需要查询的价格数据的日期。

    返回:
    返回一个列表，包含在指定日期涨停的股票代码。
    """
    # 将日期转换为格式为'YYYYMMDD'的字符串，以便在查询中使用
    date_str = date.strftime('%Y%m%d')

    # 获取指定股票列表在指定日期的每日高限价、收盘价和是否开盘的信息
    df = get_price(security=stock_list, end_date=date_str, frequency='1d', fields=['high_limit', 'close', 'is_open'], count=1)

    # 筛选出当天开盘且收盘价等于高限价的股票，即涨停的股票
    df = df[(df['is_open'] == 1) & (df['close'] == df['high_limit'])]

    # 返回涨停股票的代码列表
    return df['code'].tolist()



def get_ever_hl_stock(stock_list, date, is_break=False):
    """
    获取在指定日期达到过历史高点的股票列表。

    参数:
    stock_list: 股票代码列表，用于查询的股票池。
    date: 指定的查询日期。
    is_break: 布尔值，如果为True，则只返回在指定日期达到过历史高点但收盘价不等于历史高点且开盘的股票，默认为False。

    返回:
    股票代码列表，这些股票在指定日期达到过历史高点。
    """
    # 将日期对象转换为字符串格式，以便后续使用
    date_str = date.strftime('%Y%m%d')

    # 获取指定股票列表在指定日期的市场价格数据，包括最高价、收盘价、是否开盘等信息
    df = get_price(security=stock_list, end_date=date_str, frequency='1d', fields=['high_limit', 'high', 'close', 'is_open'], count=1)

    # 筛选出达到过历史高点的股票
    df = df[df['high'] == df['high_limit']]

    # 如果设置了is_break参数为True，则进一步筛选出那些收盘价不等于历史高点且开盘的股票
    if is_break:
        df = df[(df['close'] != df['high_limit']) & (df['is_open'] == 1)]

    # 返回筛选后的股票代码列表
    return df['code'].tolist()



# 每日初始股票池
def prepare_stock_list(date):
    """
    准备股票列表。

    根据指定日期，获取所有A股股票，并过滤掉不符合条件的股票，如ST、暂停上市、退市的股票以及新股。

    参数:
    date (datetime): 指定的日期。

    返回:
    list: 经过滤的股票列表。
    """
    # 将日期格式化为YYYYMMDD字符串，以便在数据库查询中使用
    date_str = date.strftime('%Y%m%d')

    # 获取指定日期的所有A股股票列表
    initial_list = get_Ashares(date=date_str)

    # 过滤掉ST、暂停上市、退市的股票
    initial_list = filter_stock_by_status(initial_list, filter_type=["ST", "HALT", "DELISTING"], query_date=date_str)

    # 过滤掉新股
    initial_list = filter_new_stock(initial_list, date)

    # 返回过滤后的股票列表
    return initial_list



def filter_new_stock(stock_list, date):
    """
    过滤次新股：排除上市不足250天的股票。
    """
    return [stock for stock in stock_list if not (date - datetime.strptime(get_stock_info(stock, field=['listed_date'])[stock]['listed_date'], '%Y-%m-%d').date() < timedelta(days=50))]


# 上涨时，未放量 rising on low volume    
def rise_low_volume(s, context):
    hist = get_history(count=106, frequency='1d', field=['high', 'volume'], security_list=s)
    high_prices = hist['high'][:102]
    prev_high = high_prices[-1]
    zyts_0 = next((i - 1 for i, high in enumerate(high_prices[-3::-1], 2) if high >= prev_high), 100)
    zyts = zyts_0 + 5
    if hist['volume'][-1] <= max(hist['volume'][-zyts:-1]) * 0.9:
        return True
    return False


# 计算股票处于一段时间内相对位置
def get_relative_position_df(stock_list, date, watch_days):
    if len(stock_list) != 0:
        result = pd.DataFrame()
        result = pd.DataFrame(index=stock_list)
        result.index.name = 'code'

        for sec in stock_list:
            df = get_price(security=sec, end_date=date, frequency='1d', fields=['high', 'close', 'low'], count=watch_days + 1)
            high = df['high'].shift(1).rolling(window=watch_days).max()[-1]
            low = df['low'].shift(1).rolling(window=watch_days).min()[-1]
            close = df['close'][-1]
            result.at[sec, 'rp'] = (close - low) / (high - low)
        return result

    else:
        return pd.DataFrame(columns=['rp'])


def handle_data(context, data):
    pass
